# -*- coding: utf-8 -*-
# author： woldier wong
# datetime： 3/12/24 3:15 AM
# ide： PyCharm
"""
Implements the Generalized R-CNN framework
"""

import warnings
from collections import OrderedDict
from typing import Dict, List, Optional, Tuple, Union

import torch
from torch import nn, Tensor


def _check_degenerate_boxes(targets: Optional[List[Dict[str, Tensor]]]):
    """
    检查每一个box 的x1,y1 和 x2,y2的值是否是合理的, 是否会出现 x1> x2和y1>y2的情况
    :param targets:
    :return:
    """
    if targets is not None:  # 如果traget不是None , 说明是训练模式, 这里面存的是transform变换后的bboxes
        for target_idx, target in enumerate(targets):  # 遍历
            boxes = target["boxes"]  # 得到一个bbox
            degenerate_boxes = boxes[:, 2:] <= boxes[:, :2]
            if degenerate_boxes.any():
                # print the first degenerate box
                bb_idx = torch.where(degenerate_boxes.any(dim=1))[0][0]
                degen_bb: List[float] = boxes[bb_idx].tolist()
                torch._assert(
                    False,
                    "All bounding boxes should have positive height and width."
                    f" Found invalid box {degen_bb} for target at index {target_idx}.",
                )


class GeneralizedRCNN(nn.Module):
    """
    Main class for Generalized R-CNN.

    Args:
        backbone (nn.Module):  骨干网络
        rpn (nn.Module):  区域建议网络
        roi_heads (nn.Module): takes the features + the proposals from the RPN and computes
            detections / masks from it.  兴趣池化 + 分类与回归多头
        transform (nn.Module): performs the data transformation from the inputs to feed into
            the model  将输入转换为可以作为网络输入的形式
    """

    def __init__(self, backbone: nn.Module, rpn: nn.Module, roi_heads: nn.Module, transform: nn.Module) -> None:
        super().__init__()
        self.transform = transform
        self.backbone = backbone
        self.rpn = rpn
        self.roi_heads = roi_heads
        # used only on torchscript mode
        self._has_warned = False

    @torch.jit.unused
    def eager_outputs(self, losses, detections):
        # type: (Dict[str, Tensor], List[Dict[str, Tensor]]) -> Union[Dict[str, Tensor], List[Dict[str, Tensor]]]
        """
        本方法的作用是, 根据当前模型所处的模式是训练模式还是推理模式, 动态的返回需要的值
        :param losses:
        :param detections:
        :return:
        """
        if self.training:
            return losses

        return detections

    def forward(self, images, targets=None):
        # type: (List[Tensor], Optional[List[Dict[str, Tensor]]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]
        """
        Args:
            images (list[Tensor]): images to be processed
            targets (list[Dict[str, Tensor]]): ground-truth boxes present in the image (optional)

        Returns:
            result (list[BoxList] or dict[Tensor]): the output from the model.
                During training, it returns a dict[Tensor] which contains the losses.  训练阶段返回损失
                During testing, it returns list[BoxList] contains additional fields   推理阶段返回 打分, 标签, 遮罩等
                like `scores`, `labels` and `mask` (for Mask R-CNN models).

        """
        if self.training:  # 如果是训练阶段, 判断是否给出的targets
            if targets is None:
                assert False, "targets should not be none when in training mode"
            else:  # 如果有targets 进一步判断其格式
                for target in targets:
                    boxes = target["boxes"]
                    if isinstance(boxes, torch.Tensor):
                        assert len(boxes.shape) == 2 and boxes.shape[-1] == 4, \
                            f"Expected target boxes to be a tensor of shape [N, 4], got {boxes.shape}."
                    else:
                        assert False, f"Expected target boxes to be of type Tensor, got {type(boxes)}."
        # 定义一个初始image size 的数组, 数组中的元素是一个长度为2的元组, 表示图片的h w
        original_image_sizes: List[Tuple[int, int]] = []
        for img in images:
            val = img.shape[-2:]  # 得到image 的后两个维度 得到h w
            # 判断 是否有两个维度
            assert len(val) == 2, \
                f"expecting the last two dimensions of the Tensor to be H and W instead got {img.shape[-2:]}"
            original_image_sizes.append((val[0], val[1]))
        # 调用transform 模块, 将imgae 和 bboxes 变换
        images, targets = self.transform(images, targets)

        # Check for degenerate boxes
        # TODO: Move this to a function
        _check_degenerate_boxes(targets)
        # 将image 送入backbone 抽取特征
        features = self.backbone(images.tensors)
        if isinstance(features, torch.Tensor):
            features = OrderedDict([("0", features)])

        # 送入RPN 得到候选框, 以及在train时计算PRN loos
        proposals, proposal_losses = self.rpn(images, features, targets)
        # 特征, 候选区,图片大小,targets 送入roi_head 得到 检测结果和 检测损失
        detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)
        # 后置处理, 将预测结果(bbox)转换为原始的图片尺寸
        detections = self.transform.postprocess(detections, images.image_sizes,
                                                original_image_sizes)  # type: ignore[operator]

        losses = {}
        losses.update(detector_losses)
        losses.update(proposal_losses)

        if torch.jit.is_scripting():
            if not self._has_warned:
                warnings.warn("RCNN always returns a (Losses, Detections) tuple in scripting")
                self._has_warned = True
            return losses, detections
        else:
            return self.eager_outputs(losses, detections)
